import numpy as np
import matplotlib.pyplot as plt
import math
import copy

class DescartesPoint:
    def __init__(self, x, y):
        self.X = x
        self.Y = y

    def toImagePoint(self, image_width, image_height):
        row = image_height - self.Y
        col = self.X
        return ImagePoint(row, col)

class ImagePoint:
    def __init__(self, row, col):
        self.Row = row
        self.Col = col

    def toDescartes(self, image_width, image_height):
        y = image_height - self.Row
        x = self.Col
        return DescartesPoint(x, y)
    
class LineSegmentFitForImage:
    def __init__(self, image_width, image_height):
        self.ImageWidth = image_width
        self.ImageHeight = image_height
    def fitLineSegmentForDescartesPoints(self, descartes_points):
        """最小二乘拟合直线"""
        sumx = 0
        sumy = 0
        sumxy = 0
        sumx2 = 0
        n = len(descartes_points)


        for i in range(n):
            x = descartes_points[i].X
            y = descartes_points[i].Y
            sumx += x
            sumy += y
            sumx2 += x**2
            sumxy += x * y
        
        molecule = (n * sumxy - sumx * sumy) # 分子
        denominator = (sumx**2 - n * sumx2)  # 分母
        if denominator == 0:
            # 垂直的线
            miny = 10000
            maxy = -10000
            for i in range(n):
                y = descartes_points[i].Y
                if y > maxy:
                    maxy = y
                elif y < miny:
                    miny = y
            return (DescartesPoint(descartes_points[0].X, maxy), DescartesPoint(descartes_points[0].X, miny))
        else:
            a = -(molecule / denominator)
            b = (sumy - sumx * a) / n

            if a == 0:
                # 水平线
                minx = 10000
                maxx = -10000
                for i in range(n):
                    x = descartes_points[i].X
                    if x > maxx:
                        maxx = x
                    elif x < minx:
                        minx = x
                return (DescartesPoint(minx, descartes_points[0].Y), DescartesPoint(maxx, descartes_points[0].Y))
            else:
                # y = ax + b
                # 计算y轴交点
                cross_to_y = (0, b)
                zeta = math.atan(a)
                e_vector_line = (math.cos(zeta), math.sin(zeta))
                min_dist = 10000
                max_dist = -10000
                for i in range(n):
                    x = descartes_points[i].X - cross_to_y[0]
                    y = descartes_points[i].Y - cross_to_y[1]
                    dist = e_vector_line[0]*x + e_vector_line[1]*y
                    if dist < min_dist:
                        min_dist = dist
                    if dist > max_dist:
                        max_dist = dist

                p1 = DescartesPoint(cross_to_y[0]+e_vector_line[0]*min_dist, cross_to_y[1]+e_vector_line[1]*min_dist)
                p2 = DescartesPoint(cross_to_y[0]+e_vector_line[0]*max_dist, cross_to_y[1]+e_vector_line[1]*max_dist)

                if p1.X < p2.X:
                    return (p1, p2)
                else:
                    return (p2, p1) 
    def fitLineSegmentForDescartesPoints2(self, descartes_points):
        '''整体最小二乘'''
        sumt = 0
        sums = 0
        sums2 = 0
        sumst = 0
        sumt2 = 0
        n = len(descartes_points)
        for i in range(n):
            x = descartes_points[i].X
            y = descartes_points[i].Y
            sumt += x
            sums += y
            sums2 += y ** 2
            sumst += x * y
            sumt2 += x ** 2
        
        avg_s = sums / n
        avg_t = sumt / n
        lss = sums2 - (sums ** 2) / n
        lst = sumst - (sums * sumt) / n
        ltt = sumt2 - (sumt ** 2) / n


        if lst == 0:
            # 垂直的线
            miny = 10000
            maxy = -10000
            for i in range(n):
                y = descartes_points[i].Y
                if y > maxy:
                    maxy = y
                elif y < miny:
                    miny = y
            return (DescartesPoint(descartes_points[0].X, maxy), DescartesPoint(descartes_points[0].X, miny))
        else:
            _b = lss-ltt
            c1 = (_b + math.sqrt( (_b**2)+4 * (lst ** 2) )) / (2 * lst)
            if c1 == 0:
                 # 水平线
                minx = 10000
                maxx = -10000
                for i in range(n):
                    x = descartes_points[i].X
                    if x > maxx:
                        maxx = x
                    elif x < minx:
                        minx = x
                return (DescartesPoint(minx, descartes_points[0].Y), DescartesPoint(maxx, descartes_points[0].Y))
            else:
                c2 = avg_s - c1 * avg_t
                # y = C1 x + C2
                # 计算y轴交点
                cross_to_y = (0, c2)
                cross_to_x = (-c2/c1, 0)
                vector_line = (-c2/c1, -c2)
                len_of_vector = math.sqrt(vector_line[0]**2 + vector_line[1]**2)
                e_vector_line = (vector_line[0]/len_of_vector, vector_line[1]/len_of_vector)
                min_dist = 1000000000
                max_dist = -1000000000
                for i in range(n):
                    x = descartes_points[i].X - cross_to_y[0]
                    y = descartes_points[i].Y - cross_to_y[1]
                    dist = e_vector_line[0]*x + e_vector_line[1]*y
                    if dist < min_dist:
                        min_dist = dist
                    if dist > max_dist:
                        max_dist = dist

                p1 = DescartesPoint(cross_to_y[0]+e_vector_line[0]*min_dist, cross_to_y[1]+e_vector_line[1]*min_dist)
                p2 = DescartesPoint(cross_to_y[0]+e_vector_line[0]*max_dist, cross_to_y[1]+e_vector_line[1]*max_dist)

                if p1.X < p2.X:
                    return (p1, p2)
                else:
                    return (p2, p1) 





    def fitLineSegmentForImagePoints(self, image_points):
        # 全部转换到笛卡尔坐标
        descartesPoints = []
        for image_point in image_points:
            descartesPoints.append(image_point.toDescartes(self.ImageWidth, self.ImageHeight))

        (p1,p2) = self.fitLineSegmentForDescartesPoints(descartesPoints)
        return (p1.toImagePoint(self.ImageWidth, self.ImageHeight), p2.toImagePoint(self.ImageWidth, self.ImageHeight))

    
def test_fit_line():
    fig = plt.figure()
    ax1 = fig.add_subplot(111)
    ax1.axis('equal')
    ax1.grid()

    image_width = 2000
    image_height = 1600
    # image_points = [
    #     ImagePoint(701.58, 999.578),
    #     ImagePoint(756.591, 1005.0),
    #     ImagePoint(799.665, 1010.5),
    #     ImagePoint(848.654, 999.9),
    #     ImagePoint(897.977, 1001.43)
    # ]

    image_points = [
        ImagePoint(1500, 100),
        ImagePoint(1400, 500),
        ImagePoint(800, 1000),
        ImagePoint(600, 1100),
        ImagePoint(550, 1400)
    ]
    # image_points = [
    #     ImagePoint(820.264, 902.274),
    #     ImagePoint(799.968, 951.399),
    #     ImagePoint(750.665, 1020.5),
    #     ImagePoint(799.406, 1049.57),
    #     ImagePoint(780.132, 1098.69)
    # ]

    descartesPoints = [
        image_points[0].toDescartes(image_width, image_height),
        image_points[1].toDescartes(image_width, image_height),
        image_points[2].toDescartes(image_width, image_height),
        image_points[3].toDescartes(image_width, image_height),
        image_points[4].toDescartes(image_width, image_height)
    ]
    _x = [p.X for p in descartesPoints]
    _y = [p.Y for p in descartesPoints]
    _np_x = np.array(_x)
    _np_y = np.array(_y)
    ax1.plot(_np_x, _np_y, color = 'green')

    line_fit = LineSegmentFitForImage(image_width, image_height)
    (p1, p2) = line_fit.fitLineSegmentForDescartesPoints2(descartesPoints)
    ax1.plot(np.array([p1.X, p2.X]), np.array([p1.Y, p2.Y]), color = 'red')
    plt.show()


if __name__ == "__main__":
        test_fit_line()